from torch import nn
import torch

#创建模型类
class Net(nn.Module):
    #定义参数
    def __init__(self):
        super(Net, self).__init__()

        self.w = nn.Parameter(torch.randn(2,1)) #nn.Prameter定义参数w,并初始化
        self.b = nn.Parameter(torch.zeros(1,1))
    #前向传播：forward
    def forward(self, x):
        y = x@self.w + self.b
        return y

net = Net()

from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter('./data/tensorboard')  #创建summary写对象
writer.add_graph(net,input_to_model = torch.rand(10,2))#add_graph添加网络模型图;input_to_model模型数据数据
writer.close()  #关闭写对象

#torch内直接启动tensorboard
from tensorboard import notebook
notebook.list()

notebook.start("--logdir ./data/tensorboard")

# tensorboard --logdir 路径

